# Monte Carlo Tree Search module

import math
import game

class MCTSTree (game.TreeNode):

    conf = 0.1
    threshold = 10
    bias = 1.0
    
    def __init__(self, game):

        self.wins  = 0.0
        self.plays = 0.0
        self.game  = game
        self.child = None

    def __err(self, move):

        err = MCTSTree.conf
        err = err * math.log(self.plays + 1)
        err = err * math.sqrt(1.0 / self.child[move].plays)
        return err

    def __wr(self, move):

        return (self.child[move].wins + MCTSTree.bias) / (self.child[move].plays + MCTSTree.bias * 2)

    def ucb(self, move):

        if self.child == None:
            return 1.0
        elif move not in self.child:
            return 1.0
        else:
            return self.__wr(move) + self.__err(move)

    def winrate(self, move):
        
        if self.child == None:
            return 0.0
        elif move not in self.child:
            return 0.0
        else:
            return self.__wr(move)

    def playout(self):

        if self.child is not None:
            
            move = self.choose_ucb()
            if not move in self.child:
                self.child[move] = MCTSTree(self.game.move(move))
            
            score = self.child[move].playout()
            self.plays = self.plays + 1.0
            self.wins  = self.wins + score[self.game.turn()]

            return score

        else:

            score = self.game.score()
            self.plays = self.plays + 1.0
            self.wins  = self.wins + score[self.game.turn()]

            if self.wins > MCTSTree.threshold:
                self.child = {}

            return score

    def choose_ucb(self):
        
        best = None
        maximum = 0.0
        for m in self.game.valid():
            if self.ucb(m) > maximum:
                best = m
                maximum = self.ucb(m)

        return best

    def choose_winrate(self):
        
        best = None
        maximum = 0.0
        for m in self.game.valid():
            if self.winrate(m) > maximum:
                best = m
                maximum = self.winrate(m)

        return best

    def choose(self):
        
        return self.choose_winrate()

    def get_score(self):
       
        return (self.wins + MCTSTree.bias) / (self.plays + MCTSTree.bias * 2)

    def get_child(self, move):
        
        if self.child != None:
            if move in self.child:
                return self.child[move]
            else:
                return None
        else:
            return None
